import torch
from torch import nn
from training.models.resnet import BasicBlock

__all__ = ['discriminativecolorspace']


class ColorProjectModel(nn.Module):

    def __init__(self, inplanes, planes):
        super(ColorProjectModel, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, 3, 1, 1, 1)
        self.act1 = nn.Sigmoid()

        self.conv2 = nn.Conv2d(inplanes, 3, 1, 1, 1)
        self.act2 = nn.Sigmoid()


    def forward(self, x):
        x = self.conv1(x)
        x = self.act1(x)

        x = self.conv2(x)
        x = self.act2(x)

        return x


class DiscriminativeColorSpace(nn.Module):

    def __init__(self, feature_extractor, in_channels, num_classes, pretrained=False):
        super(DiscriminativeColorSpace, self).__init__()
        self.compact = ColorProjectModel(in_channels, 3)
        self.features = feature_extractor(pretrained=pretrained, num_classes=num_classes)

    def forward(self, x):
        x = self.compact(x)
        x = self.features(x)
        return x


def discriminativecolorspace(feature_extractor, in_channels=3, num_classes=1000, pretrained=False):
    model = DiscriminativeColorSpace(feature_extractor, in_channels=in_channels, num_classes=num_classes, pretrained=pretrained)
    model.default_cfg = model.features.default_cfg
    return model
